#!/usr/bin/env python
# encoding: utf-8
"""
TwoPartyLegislature.py

Created by Drew Conway <agc282@nyu.edu> on 2008-10-07.
Copyright (c) 2008. All rights reserved, see GNU LGPL.

PyVote is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.


Version = 0.1a
"""

import sys
import os
import unittest
from scipy import stats,random,array,median,mean,std
from math import sqrt,pi,e
import pylab as P

class TwoPartyLegislature:
# Object where agents are created and exist.  Agents can fit a rough random bi-modal distribution when
# ditribution is set to 'None', or a distribution can be provided by the user.  This distribution must
# be a list containing 10 ints that sum to 'num_agents'.  This construction is used to fix the distribution 
# of ideologies within the legislautre.
	def __init__(self,num_agents,distribution=None,tol=None):
		# Only the number of agents and a possible ideological distribution are provided outside the class
		
		# Class variables provided inside the class
		self.agents=[]	# Agents in legislature
		self.bills=[]	# Bills voted on by legislature
		self.num_liberals=0	
		self.num_conservatives=0
		self.ideologies=[]	# List of agent ideologies
		self.tolerances=[]	# List of adjusted agent tolerances
		
		
		if distribution is None:
		# When no ideological distribution is provided, fill out legislature with agents and do
		# appropriate record keeping.
			for a in range(0,num_agents):
				A=TwoPartyAgent(tolerance_function=tol)
				A.set_id(a)
				if A.get_party()<1:
					self.num_liberals=self.num_liberals+1
					self.ideologies.append(A.get_ideology())
					self.tolerances.append(A.get_tolerance())
					self.agents.append(A)
				else:
					self.num_conservatives=self.num_conservatives+1
					self.ideologies.append(A.get_ideology())
					self.tolerances.append(A.get_tolerance())
					self.agents.append(A)
		else:
		# If distribution is provided, first check to make sure the breakdown equals the total number of agents
		# inside the legislature
			if sum(distribution) != num_agents:
				raise IndexError('Distribution count must equal the number of agents')
			else:
				agent_count=0
				for i in range(0,len(distribution)):
					agents_in_bin_count=0
					while agents_in_bin_count<distribution[i]:
					# Next, create agents that fit the given ideolgical distribution, and perform appropraite record keeping
						A=TwoPartyAgent(tolerance_function=tol)
						if (A.ideology-(float(i)/len(distribution)))>0.0 and (A.ideology-(float(i)/len(distribution)))<0.1:
						# this test makes sure the agent created fits the ideological bin being filled
							if A.get_party()<1:
								self.num_liberals=self.num_liberals+1
								self.ideologies.append(A.get_ideology())
								self.tolerances.append(A.get_tolerance())
								A.set_id(agent_count)
								self.agents.append(A)
								agents_in_bin_count=agents_in_bin_count+1
								agent_count=agent_count+1
							else:
								self.num_conservatives=self.num_conservatives+1
								self.ideologies.append(A.get_ideology())
								self.tolerances.append(A.get_tolerance())
								A.set_id(agent_count)
								self.agents.append(A)
								agents_in_bin_count=agents_in_bin_count+1
								agent_count=agent_count+1
					
		
		
	def size(self):
	# return legislature size
		return len(self.agents)
		
	def legislators(self):
	# return a list of agents in legislature
		return self.agents
	
	def get_num_liberals(self):
		return self.num_liberals
		
	def get_num_conservatives(self):
		return self.num_conservatives
		
	def add_bill(self, bill):
		self.bills.append(bill)
		
	def get_all_bills(self):
		return self.bills
		
	def get_passed_bills(self):
		passed=[]
		for b in self.bills:
			if b.bill_pass:
				passed.append(b)
		return passed
		
	def get_agent(self, agent_id):
	# return a specific agent from legislature
		return self.agents[agent_id]
		
	def find_liberals(self):
	# return a list of all liberal agents in legislature
		liberals_index=[]
		for i in range(0,len(self.agents)):
			if self.agents[i].party<1:
				liberals_index.append(i)
		return liberals_index
		
	def find_conservatives(self):
	# return a list of all conservative agents in legislature
		conservatives_index=[]
		for i in range(0,len(self.agents)):
			if self.agents[i].party>0:
				conservatives_index.append(i)
		return conservatives_index

	
	def info(self):
	# prints basic information about legislature
		print 'Number of liberals = '+str(self.num_liberals)
		print 'Number of conservatives = '+str(self.num_conservatives)
		print 'Median ideology = '+str(median(self.ideologies))
		print 'Standard deviation of tolerances = '+str(std(self.tolerances))
		print ''
		
	def show_ideolody_distribution(self,num_bins=10):
	# displays a histogram of ideological distribution in legislature
		P.figure(1)
		h,r=P.histogram(self.ideologies,num_bins)
		P.bar(r,h,width=0.05)
		P.title('Distrobution of ideologies')
		P.xlabel('Ideologies')
		P.ylabel('Number of agents')
		
	def show_tolerance_distribution(self,num_bins=10):
	# displays a histogram of tolerance distribution in legislature
		P.figure(2)
		h,r=P.histogram(self.tolerances,num_bins)
		P.bar(r,h,width=0.05)
		P.title('Distrobution of tolerances')
		P.xlabel('Tolerances')
		P.ylabel('Number of agents')

	def show_tolerance_scatter(self):
	# displays a scatter plot of ideologies (X) by tolerances (Y)
		P.figure(3)
		P.scatter(self.ideologies,self.tolerances)
		P.title('Real ideology X Tolerance')
		P.xlabel('Real ideology')
		P.ylabel('Tolerance')

class Committee:
# Returns a committee with a given size from a given TwoPartyLegislature class capable of producing a bill
	def __init__(self,legislature,size):
		# Class variables provided outside the class
		self.legislature_size=legislature.size()
		# committe size is given by user
		
		# Class variables provided inside the class
		self.committee=[]		
		self.committee_index=[]	# Record keeping used to prevent duplicate agents on single committee
		self.liberal_count=0	# Additional record keeping to provide majority rule
		self.conservative_count=0
		self.num_liberal_members=int	# Final records for committee balance
		self.num_conservative_members=int

		# The legislature's party balance determines the number of agents from each party on the committee 
		if legislature.get_num_liberals()>legislature.get_num_conservatives():
		# If the legislature has a liberal majority, or
			self.num_liberal_members=(size/2)+1
			self.num_conservative_members=(size/2)
		else:
		# if the legislature has a conservative majority
			self.num_liberal_members=(size/2)
			self.num_conservative_members=(size/2)+1
		
		# Now randonly select agents from the legislature w/o replacement
		while len(self.committee)<size:
			i=int(random.random_integers(0,self.legislature_size-1)) # random agent
			if legislature.get_agent(i).get_party()<1:
			# If agent is a liberal, and
				if self.liberal_count<self.num_liberal_members:
				# we have not already reached our maxiumum number of liberals, and
					if self.committee_index.count(i)<1:
					# we have not already selected this agent, then add the agent to the committee and do
					# appropriate record keeping
						self.committee.append(legislature.get_agent(i))
						self.committee_index.append(i)
						self.liberal_count=self.liberal_count+1
			else:
			# Analogous process for selecting conservative agent
				if self.conservative_count<self.num_conservative_members:
					if self.committee_index.count(i)<1:
						self.committee.append(legislature.agents[i])
						self.committee_index.append(i)
						self.conservative_count=self.conservative_count+1

	def make_bill(self,majority_rule=True):
	# Committee bill generating function, majority rule means only majority party can introduce bills, otherwise random.
	# Ideologies are always set to that of the median agent from the party introducing the bill.
		bill=Bill()
		if majority_rule:
			majority_ideologies=[]
			if self.num_liberal_members>self.num_conservative_members:
			# If committee has a liberal majority set bill party to 0, then
				bill.set_party(0)
				for m in self.committee:
				# record ideologies of all liberal members, then
					if m.get_party()<1:
						majority_ideologies.append(m.get_ideology())
			else:
			# Analogous for conservative majority
				bill.set_party(1)
				for m in self.committee:
					if m.get_party()>0:
						majority_ideologies.append(m.get_ideology())
			# set bill ideology to the median ideology of the majority members
			bill.set_ideology(median(majority_ideologies))
		else:
		# If not using majority rule, randomly select the party introducing the bill, then
		# collect ideologies of party introducing bill
			majority_ideologies=[]
			if int(random.random_integers(0,2))<1: 
			# If random party selected is liberal, set bill party, then
				bill.set_party(0)
				for m in self.committee:
				# collect liberal ideologies, and 
					if m.get_party()<1:
						majority_ideologies.append(m.get_ideology())
			else:
				bill.set_party(1)
				for m in self.committee:
					if m.get_party()>0:
						majority_ideologies.append(m.get_ideology())
			# set bill ideology to median majority agent
			bill.set_ideology(median(majority_ideologies))

		return bill
		
	def show_members(self):
	# Prints agent IDs of committee members
		member_ids=[]
		for m in self.committee:
			member_ids.append(m.get_id())
		print 'Committee members: '+str(member_ids)
		
	def get_members(self):
	# Returns list of committee agents
		return self.committee
		
	def get_num_liberals(self):
	# Return count of liberal members
		return self.num_liberals
		
	def get_num_conservatives(self):
	# Return count of conservative members
		return self.num_conservatives
		
	

class Bill:
# The object that is voted on by agents
	def __init__(self):
		# Class variables provided outside the class
		self.party=int 			# The party that introduces bill
		self.ideology=float 	# 0.0 < liberal < 0.5 < conservative < 1.0
		self.record={}			# Voting record, keys: agent object, values: binary
		self.id=int				# Bill number
		
		# Class variables provided inside the class
		self.bill_pass=False	# True when bill is passed
	
	def get_party(self):
	# Return party that introduced bill
		return self.party
	
	def get_ideology(self):
	# Return bill ideology
		return self.ideology
	
	def get_bill_num(self):
	# Return bill numbers
		return self.id
		
	def get_record(self):
	# Return agent voting record of bill
		return self.record
	
	def set_bill_num(self,num):
	# Assign bill number
		self.id=num

	def set_party(self,party_input):
	# Assign party that introduced bill
		self.party=party_input
	
	def set_ideology(self,ideology_input):
	# Assign bill ideology
		self.ideology=ideology_input
	
	def set_passage(self,pass_bool):
	# Set True if bill passed
		self.bill_pass=pass_bool
		
	def add_record(self,agent,vote):
	# Add to bill voting record, keys: TwoPartyAgent, values: binary
		self.record[agent]=vote
		
	def num_liberal_yays(self):
	# return number of liberal yay votes
		liberal_yays=0
		record_items=self.record.items()
		yay_voters=[(a) for (a,b) in record_items if b>0]
		for a in yay_voters:
			if a.get_party()<1:
				liberal_yays=liberal_yays+1
		return liberal_yays
				
	def num_conservative_yays(self):
	# return number of conservative yay votes
		conservative_yays=0
		record_items=self.record.items()
		yay_voters=[(a) for (a,b) in record_items if b>0]
		for a in yay_voters:
			if a.get_party()>0:
				conservative_yays=conservative_yays+1
		return conservative_yays
	
	def info(self):
	# Returns key bill information 
		print 'Number: '+str(self.id)
		if self.get_party()<1:
			print 'Introduced by liberal'
		else:
			print 'Party that introduced: Conservative'
		print 'Ideological score: '+str(self.ideology)
		if len(self.record)<1:
			print 'Bill has not yet been voted on'
		else:
			lib_yays=self.num_liberal_yays()
			cons_yays=self.num_conservative_yays()
			if self.bill_pass:
				print 'Bill passed with '+str(lib_yays+cons_yays)+' yay votes'
			else:
				print 'Bill did not pass with '+str(lib_yays+cons_yays)+' yay votes'
		print '--- Vote Breakout ---'
		print 'Liberal yay votes: '+str(lib_yays)
		print 'Conservative yay votes: '+str(cons_yays)


class TwoPartyAgent:
# Agent class taking either a 'liberal' or 'conservative' party and ideology
	def __init__(self,ideology_input=None,tolerance_function=None):
		# Class variables given outside the class init
		self.id=int
		self.record={}	# Voting record, keys: bill num, values: binary
		
		# Default tolerance calculated as a form of the normal distribution PDF,
		# this is an area that needs to be fleshed out in much greater detail
		# LaTex of the function: $\frac{1}{1.4\sqrt{2\pi}}e^{\frac{-(t-0.5)^{2}}{(0.5)^{2}}}$
		self.tolerance=float 	# Assigned via the lambda function below
		if tolerance_function is None:
			tolerance_function = lambda t: (1.0/(1.4*sqrt(2.0*pi)))*e**(-1.0*(t-0.5)**2.0/(0.5)**2)
		
		# Agent's ideology is given by the user (unconventional initialization)
		if ideology_input is not None:
			if ideology_input <0.0 or ideology_input>1.0:
			# Value must be between 0.0 and 1.0
				raise ValueError('Agent ideology must be 0.0 < liberal < 0.5 < conservative < 1.0')
			else:
			# In this construction, if an agent is given an ideolgy of exactly 0.5, it is liberal,
			# this could also be determined by a coin toss
				if ideology_input<=0.5:	
					self.party=0
				else:
					self.party=1
				self.ideology=ideology_input
				self.tolerance=self.ideology+tolerance_function(self.ideology)
		else:
		# Agent's ideology is not given, generate these value internally (conventional initialization)
			self.ideology=float # 0.0 < liberal < 0.5 < conservative < 1.0
		
			# Class variables given inside the class init
			self.party=int(random.randint(0,2))
			ideology_vector=random.sample(100)
	
			# Liberal agent consruction
			if self.party<1:
			# Ideology is selected randomly from [0.0, 0.5]
				ideology_vector.sort()
				ideology_vector.resize(50)
				random.shuffle(ideology_vector)
				self.ideology=ideology_vector[0]
				self.tolerance=self.ideology+tolerance_function(self.ideology)
			# Conservative agent consruction
			else:
			# Ideology is selected randomly from [0.5, 1.0]	
				ideology_vector.sort()
				ideology_vector=list(ideology_vector)
				ideology_vector.reverse()
				ideology_vector=array(ideology_vector)
				ideology_vector.resize(50)
				random.shuffle(ideology_vector)
				self.ideology=ideology_vector[0]
				self.tolerance=self.ideology-tolerance_function(self.ideology)
	
	
	def get_record(self):
	# returns agent voting record
		return self.record
		
	def get_ideology(self):
	# returns agent ideology
		return self.ideology
		
	def get_tolerance(self):
	# returns agent adjusted tolerance
		return self.tolerance
		
	def get_id(self):
	# returns agent ID
		return self.id
		
	def get_party(self):
	# returns agent party (0=liberal, 1=conservative)
		return self.party
	
	def set_id(self, agent_id):
	# Assign agent ID
		if not isinstance(agent_id,int):
			raise ValueError("id's must be ints not %s"%str(type(agent_id)))
		self.id=agent_id
	
	def add_record(self, bill_num,vote):
	# Add to agent voting record
		if vote not in (0,1):
			raise ValueError("Agent vote must be 0 or 1")
		else:
			self.record[bill_num]=vote
			
	def info(self):
	# Prints key agent information
		print 'ID: '+str(self.id)
		if self.party<1:
			print 'Member of liberal party'
		else:
			print 'Member of conservative party'
		print 'Ideology: '+str(self.ideology)

		

class TwoPartyAgentTests(unittest.TestCase):

    agent=None


    def get_agent(self,ideology_input=None,tolerance_function=None):
        """Assemble a TwoPartyAgent with specified constructor arguments"""
        #Override so that party is always the same for testing
        f,random.randint=random.randint,lambda x,y:0
        agent=TwoPartyAgent(ideology_input,tolerance_function)
        random.randint=f
        return agent

    def setUp(self):
        self.agent=self.get_agent()

    def test_init_ideology_out_of_range(self):
        """ValueError should be raised when ideology is out of range"""
        self.assertRaises(ValueError,TwoPartyAgent,-.1)
        self.assertRaises(ValueError,TwoPartyAgent,1.1)

    def test_init_ideology_sets_party(self):
        """Passing in an ideology should set the appropriate party affiliation"""
        agent=self.get_agent(.1)
        self.assertEquals(0,agent.get_party())
        agent=self.get_agent(.9)
        self.assertEquals(1,agent.get_party())

    def test_init_tolerance_from_tolerancefunc_and_ideology(self):
        """A passed in tolerance function should be used to generate the tolerance level"""
        f=lambda x:x+1
        agent=self.get_agent(0.1,f)
        self.assertEquals(.1+f(.1),agent.get_tolerance())
        agent=self.get_agent(0.9,f)
        self.assertEquals(0.9+f(0.9),agent.get_tolerance())

    def test_init_tolerance_from_ideology(self):
        """Ideology specified should result in specific tolerance"""
        agent=self.get_agent(0.1)
        self.assertEquals(0.25025660149009504,agent.get_tolerance())
        
    def test_get_party(self):
        """Party should equal zero due to party def override"""
        party=self.agent.get_party()
        self.assertEquals(0,party)


    def test_get_record(self):
        """Voting record should be a dictionary"""
        record=self.agent.get_record()
        self.assert_(isinstance(record,dict))


    def test_set_id(self):
        """set_id() should modify result of get_id()"""
        self.agent.set_id(3)
        self.assertEquals(3,self.agent.get_id())

    def test_set_id_invalid_type(self):
        """set_id() should raise error if id is not int"""
        self.assertRaises(ValueError,self.agent.set_id,'foo')

    def test_add_record_invalid_vote(self):
        """Invalid vote should raise an error"""
        self.assertRaises(ValueError,self.agent.add_record,1,1.1)
        self.assertRaises(ValueError,self.agent.add_record,1,-.1)
        self.assertRaises(ValueError,self.agent.add_record,1,.5)

    def test_add_record(self):
        """Added records should be returned by get record"""
        self.agent.add_record(1,0)
        self.agent.add_record(2,1)
        self.agent.add_record(3,1)
        votedict={1:0,2:1,3:1}
        self.assertEquals(votedict,self.agent.get_record())

    def test_info_appropriate_party(self):
        """info() should display the appropriate party"""
        #Redirect print statements to StringIO object
        import StringIO
        buff=StringIO.StringIO()
        sys.stdout=buff
        self.agent.info()
        sys.stdout=sys.__stdout__
        self.assertEquals('Member of liberal party',buff.getvalue().split('\n')[1])
			


if __name__ == '__main__':
	unittest.main()

